{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import transforms, datasets\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision.models as models\n",
    "from torchsummary import summary\n",
    "from config.root_path import DATA_ROOT\n",
    "from urllib.request import urlopen\n",
    "from PIL import Image\n",
    "import timm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_path = os.path.join(DATA_ROOT,'8_class_select')\n",
    "\n",
    "data_transform = {\n",
    "        \"train\": transforms.Compose([transforms.CenterCrop(256),\n",
    "                                     transforms.Resize(152),\n",
    "                                     transforms.RandomHorizontalFlip(),\n",
    "                                     transforms.ToTensor(),\n",
    "                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                                    ]),\n",
    "        \"val\": transforms.Compose([\n",
    "                                   transforms.CenterCrop(256),\n",
    "                                   transforms.Resize(152),\n",
    "                                   transforms.ToTensor(),\n",
    "                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                                   ])\n",
    "                }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 8 dataloader workers every process\n",
      "\n"
     ]
    }
   ],
   "source": [
    "batch_size = 64\n",
    "lr = 0.0001\n",
    "nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers\n",
    "print('Using {} dataloader workers every process\\n'.format(nw))\n",
    "\n",
    "\n",
    "train_dataset = datasets.ImageFolder(os.path.join(root_path,'train'),\n",
    "                                     transform=data_transform['train'])\n",
    "\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size = batch_size,                         \n",
    "                          shuffle=True,\n",
    "                          num_workers = nw)\n",
    "\n",
    "val_dataset = datasets.ImageFolder(os.path.join(root_path,'val'),\n",
    "                                   transform=data_transform['val'])\n",
    "val_loader = DataLoader(val_dataset,\n",
    "                        batch_size = batch_size,\n",
    "                        shuffle=False,\n",
    "                        num_workers = nw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda:3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/wide_resnet50_2-9ba9bcbe.pth\" to /home/wangcheng/.cache/torch/hub/checkpoints/wide_resnet50_2-9ba9bcbe.pth\n",
      "100%|██████████| 263M/263M [00:22<00:00, 12.4MB/s] \n"
     ]
    }
   ],
   "source": [
    "net = models.wide_resnet50_2(weights='IMAGENET1K_V2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RegNet(\n",
       "  (stem): SimpleStemIN(\n",
       "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU(inplace=True)\n",
       "  )\n",
       "  (trunk_output): Sequential(\n",
       "    (block1): AnyStage(\n",
       "      (block1-0): ResBottleneckBlock(\n",
       "        (proj): Conv2dNormActivation(\n",
       "          (0): Conv2d(32, 224, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(32, 224, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(224, 224, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=2, bias=False)\n",
       "            (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(224, 8, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(8, 224, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(224, 224, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block1-1): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(224, 224, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(224, 224, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2, bias=False)\n",
       "            (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(224, 56, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(56, 224, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(224, 224, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (block2): AnyStage(\n",
       "      (block2-0): ResBottleneckBlock(\n",
       "        (proj): Conv2dNormActivation(\n",
       "          (0): Conv2d(224, 448, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(224, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=4, bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(448, 56, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(56, 448, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block2-1): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=4, bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(448, 112, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(112, 448, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block2-2): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=4, bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(448, 112, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(112, 448, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block2-3): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=4, bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(448, 112, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(112, 448, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 448, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(448, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (block3): AnyStage(\n",
       "      (block3-0): ResBottleneckBlock(\n",
       "        (proj): Conv2dNormActivation(\n",
       "          (0): Conv2d(448, 1232, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(448, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 112, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(112, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-1): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-2): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-3): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-4): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-5): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-6): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-7): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-8): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-9): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "      (block3-10): ResBottleneckBlock(\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=11, bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1232, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 1232, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 1232, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1232, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (block4): AnyStage(\n",
       "      (block4-0): ResBottleneckBlock(\n",
       "        (proj): Conv2dNormActivation(\n",
       "          (0): Conv2d(1232, 3024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(3024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (f): BottleneckTransform(\n",
       "          (a): Conv2dNormActivation(\n",
       "            (0): Conv2d(1232, 3024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(3024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (b): Conv2dNormActivation(\n",
       "            (0): Conv2d(3024, 3024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=27, bias=False)\n",
       "            (1): BatchNorm2d(3024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "          (se): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(3024, 308, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(308, 3024, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): ReLU()\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (c): Conv2dNormActivation(\n",
       "            (0): Conv2d(3024, 3024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(3024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (activation): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): Linear(in_features=3024, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.classifier[2]=nn.Linear(1024,8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvNeXt(\n",
       "  (features): Sequential(\n",
       "    (0): Conv2dNormActivation(\n",
       "      (0): Conv2d(3, 128, kernel_size=(4, 4), stride=(4, 4))\n",
       "      (1): LayerNorm2d((128,), eps=1e-06, elementwise_affine=True)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(128, 128, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=128)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=128, out_features=512, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=512, out_features=128, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.0, mode=row)\n",
       "      )\n",
       "      (1): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(128, 128, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=128)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=128, out_features=512, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=512, out_features=128, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.014285714285714285, mode=row)\n",
       "      )\n",
       "      (2): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(128, 128, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=128)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((128,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=128, out_features=512, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=512, out_features=128, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.02857142857142857, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): LayerNorm2d((128,), eps=1e-06, elementwise_affine=True)\n",
       "      (1): Conv2d(128, 256, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(256, 256, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=256)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=256, out_features=1024, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=1024, out_features=256, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.04285714285714286, mode=row)\n",
       "      )\n",
       "      (1): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(256, 256, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=256)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=256, out_features=1024, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=1024, out_features=256, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.05714285714285714, mode=row)\n",
       "      )\n",
       "      (2): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(256, 256, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=256)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((256,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=256, out_features=1024, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=1024, out_features=256, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.07142857142857142, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (4): Sequential(\n",
       "      (0): LayerNorm2d((256,), eps=1e-06, elementwise_affine=True)\n",
       "      (1): Conv2d(256, 512, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.08571428571428572, mode=row)\n",
       "      )\n",
       "      (1): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.1, mode=row)\n",
       "      )\n",
       "      (2): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.11428571428571428, mode=row)\n",
       "      )\n",
       "      (3): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.12857142857142856, mode=row)\n",
       "      )\n",
       "      (4): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.14285714285714285, mode=row)\n",
       "      )\n",
       "      (5): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.15714285714285714, mode=row)\n",
       "      )\n",
       "      (6): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.17142857142857143, mode=row)\n",
       "      )\n",
       "      (7): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.18571428571428572, mode=row)\n",
       "      )\n",
       "      (8): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.2, mode=row)\n",
       "      )\n",
       "      (9): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.21428571428571427, mode=row)\n",
       "      )\n",
       "      (10): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.22857142857142856, mode=row)\n",
       "      )\n",
       "      (11): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.24285714285714285, mode=row)\n",
       "      )\n",
       "      (12): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.2571428571428571, mode=row)\n",
       "      )\n",
       "      (13): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.2714285714285714, mode=row)\n",
       "      )\n",
       "      (14): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.2857142857142857, mode=row)\n",
       "      )\n",
       "      (15): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.3, mode=row)\n",
       "      )\n",
       "      (16): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.3142857142857143, mode=row)\n",
       "      )\n",
       "      (17): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.32857142857142857, mode=row)\n",
       "      )\n",
       "      (18): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.34285714285714286, mode=row)\n",
       "      )\n",
       "      (19): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.35714285714285715, mode=row)\n",
       "      )\n",
       "      (20): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.37142857142857144, mode=row)\n",
       "      )\n",
       "      (21): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.38571428571428573, mode=row)\n",
       "      )\n",
       "      (22): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.4, mode=row)\n",
       "      )\n",
       "      (23): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.4142857142857143, mode=row)\n",
       "      )\n",
       "      (24): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.42857142857142855, mode=row)\n",
       "      )\n",
       "      (25): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.44285714285714284, mode=row)\n",
       "      )\n",
       "      (26): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(512, 512, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=512)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.45714285714285713, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (6): Sequential(\n",
       "      (0): LayerNorm2d((512,), eps=1e-06, elementwise_affine=True)\n",
       "      (1): Conv2d(512, 1024, kernel_size=(2, 2), stride=(2, 2))\n",
       "    )\n",
       "    (7): Sequential(\n",
       "      (0): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(1024, 1024, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=1024)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.4714285714285714, mode=row)\n",
       "      )\n",
       "      (1): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(1024, 1024, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=1024)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.4857142857142857, mode=row)\n",
       "      )\n",
       "      (2): CNBlock(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2d(1024, 1024, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), groups=1024)\n",
       "          (1): Permute()\n",
       "          (2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n",
       "          (3): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (4): GELU(approximate='none')\n",
       "          (5): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (6): Permute()\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.5, mode=row)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "  (classifier): Sequential(\n",
       "    (0): LayerNorm2d((1024,), eps=1e-06, elementwise_affine=True)\n",
       "    (1): Flatten(start_dim=1, end_dim=-1)\n",
       "    (2): Linear(in_features=1024, out_features=8, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.classifier = nn.Sequential(\n",
    "    dropout,\n",
    "    nn.Linear(in_features, 8)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EfficientNet(\n",
       "  (features): Sequential(\n",
       "    (0): Conv2dNormActivation(\n",
       "      (0): Conv2d(3, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (1): BatchNorm2d(24, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): SiLU(inplace=True)\n",
       "    )\n",
       "    (1): Sequential(\n",
       "      (0): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(24, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.0, mode=row)\n",
       "      )\n",
       "      (1): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(24, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(24, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.005, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(24, 96, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(96, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(96, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(48, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.01, mode=row)\n",
       "      )\n",
       "      (1): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(48, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(192, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(192, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(48, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.015000000000000003, mode=row)\n",
       "      )\n",
       "      (2): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(48, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(192, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(192, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(48, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.02, mode=row)\n",
       "      )\n",
       "      (3): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(48, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(192, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(192, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(48, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.025, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(48, 192, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(192, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.030000000000000006, mode=row)\n",
       "      )\n",
       "      (1): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(64, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.035, mode=row)\n",
       "      )\n",
       "      (2): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(64, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.04, mode=row)\n",
       "      )\n",
       "      (3): FusedMBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(64, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.045, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (4): Sequential(\n",
       "      (0): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=256, bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.05, mode=row)\n",
       "      )\n",
       "      (1): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.05500000000000001, mode=row)\n",
       "      )\n",
       "      (2): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.06000000000000001, mode=row)\n",
       "      )\n",
       "      (3): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.065, mode=row)\n",
       "      )\n",
       "      (4): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.07, mode=row)\n",
       "      )\n",
       "      (5): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512, bias=False)\n",
       "            (1): BatchNorm2d(512, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.075, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(768, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n",
       "            (1): BatchNorm2d(768, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(768, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(32, 768, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(768, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.08, mode=row)\n",
       "      )\n",
       "      (1): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.085, mode=row)\n",
       "      )\n",
       "      (2): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.09, mode=row)\n",
       "      )\n",
       "      (3): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.095, mode=row)\n",
       "      )\n",
       "      (4): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.1, mode=row)\n",
       "      )\n",
       "      (5): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.10500000000000001, mode=row)\n",
       "      )\n",
       "      (6): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.11000000000000001, mode=row)\n",
       "      )\n",
       "      (7): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.11500000000000002, mode=row)\n",
       "      )\n",
       "      (8): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(160, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.12000000000000002, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (6): Sequential(\n",
       "      (0): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=960, bias=False)\n",
       "            (1): BatchNorm2d(960, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(960, 40, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(40, 960, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(960, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.125, mode=row)\n",
       "      )\n",
       "      (1): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.13, mode=row)\n",
       "      )\n",
       "      (2): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.135, mode=row)\n",
       "      )\n",
       "      (3): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.14, mode=row)\n",
       "      )\n",
       "      (4): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.14500000000000002, mode=row)\n",
       "      )\n",
       "      (5): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.15, mode=row)\n",
       "      )\n",
       "      (6): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.155, mode=row)\n",
       "      )\n",
       "      (7): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.16, mode=row)\n",
       "      )\n",
       "      (8): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.165, mode=row)\n",
       "      )\n",
       "      (9): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.17, mode=row)\n",
       "      )\n",
       "      (10): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.175, mode=row)\n",
       "      )\n",
       "      (11): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.18, mode=row)\n",
       "      )\n",
       "      (12): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.185, mode=row)\n",
       "      )\n",
       "      (13): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.19, mode=row)\n",
       "      )\n",
       "      (14): MBConv(\n",
       "        (block): Sequential(\n",
       "          (0): Conv2dNormActivation(\n",
       "            (0): Conv2d(256, 1536, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (1): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 1536, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=1536, bias=False)\n",
       "            (1): BatchNorm2d(1536, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): SiLU(inplace=True)\n",
       "          )\n",
       "          (2): SqueezeExcitation(\n",
       "            (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "            (fc1): Conv2d(1536, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (fc2): Conv2d(64, 1536, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (activation): SiLU(inplace=True)\n",
       "            (scale_activation): Sigmoid()\n",
       "          )\n",
       "          (3): Conv2dNormActivation(\n",
       "            (0): Conv2d(1536, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (stochastic_depth): StochasticDepth(p=0.195, mode=row)\n",
       "      )\n",
       "    )\n",
       "    (7): Conv2dNormActivation(\n",
       "      (0): Conv2d(256, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (1): BatchNorm2d(1280, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): SiLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "  (classifier): Sequential(\n",
       "    (0): Dropout(p=0.2, inplace=True)\n",
       "    (1): Linear(in_features=1280, out_features=8, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.to(device)\n",
    "for param in net.parameters():\n",
    "    print(param.device)\n",
    "    break  # 只查看第一个参数的设备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_funtion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(net.parameters(),lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(10):\n",
    "    net.train()\n",
    "    running_loss = 0.0\n",
    "    for i,data in enumerate(train_loader):\n",
    "        inputs,labels = data\n",
    "        inputs,labels = inputs.to(device),labels.to(device)\n",
    "        print(inputs.device,labels.device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(inputs)\n",
    "        loss = loss_funtion(outputs,labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "        print(running_loss)\n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sd",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
